# _*_ coding: utf-8 _*_
# ---------------------------
# @创建作者：ZQQ
# @创建时间：2025/8/28 14:50
# @说明：
# ---------------------------
import os
import time
import zipfile
from pathlib import Path

import requests


class ModelDownloader:
    def __init__(self, url, folder, chunk_size=8192, max_retries=3):
        """
        初始化下载器实例
        :param url: 文件下载URL
        :param folder: 存储目录名
        :param chunk_size: 分块大小(字节)
        :param max_retries: 最大重试次数
        """
        # 判断当前目录是否存在
        deploy_dir = Path.cwd() / folder
        print(f'当前目录:{deploy_dir}')
        if not os.path.exists(deploy_dir):
            # 自动创建多级目录
            deploy_dir.mkdir(parents=True, exist_ok=True)

        # 生成保存路径(下载压缩包)
        self.zip_path = deploy_dir / "model.zip"
        self.temp_path = self.zip_path.with_suffix('.tmp')  # 临时文件路径
        self.extract_path = deploy_dir  # 解压路径
        self.url = url
        self.chunk_size = chunk_size
        self.max_retries = max_retries

    def _extract_zip(self) -> str:
        """
        解压zip文件
        :return: 解压后的model.joblib文件路径
        """
        try:
            with zipfile.ZipFile(self.zip_path, 'r') as zip_ref:
                zip_ref.extractall(self.extract_path)
                print('✅ 模型压缩包解压完成！')
            
            # 删除压缩包以节省空间
            self.zip_path.unlink()
            print('🗑️ 压缩包已删除')
            
            # 查找解压后的model.joblib文件
            model_path = self.extract_path / "model.joblib"
            if model_path.exists():
                return str(model_path.resolve())
            else:
                # 如果在子目录中，尝试查找
                for root, dirs, files in os.walk(self.extract_path):
                    if "model.joblib" in files:
                        return str(Path(root) / "model.joblib")
                print('❌ 未找到model.joblib文件')
                return ""
        except Exception as e:
            print(f'❌ 解压失败: {e}')
            return ""

    @property
    def download(self) -> str:
        """
        执行下载逻辑（自动跳过已下载文件）
        :return: 解压后的model.joblib文件绝对路径（失败返回空字符串）
        """
        # 检查是否已经解压过
        model_path = self.extract_path / "model.joblib"
        if model_path.exists():
            print(f'✅ 模型文件已存在: {model_path}')
            return str(model_path.resolve())

        # 检查压缩包是否已存在
        if self.zip_path.exists():
            print(f'✅ 压缩包已存在，直接解压: {self.zip_path}')
            return self._extract_zip()

        # 断点续传准备
        existing_size = self.temp_path.stat().st_size if self.temp_path.exists() else 0

        # 重试下载逻辑（指数退避）
        for attempt in range(self.max_retries):
            try:
                headers = {'Range': f'bytes={existing_size}-'} if existing_size else {}

                with requests.get(self.url, headers=headers, stream=True, timeout=10) as r:
                    r.raise_for_status()

                    # 流式写入文件
                    mode = 'ab' if existing_size else 'wb'
                    with open(self.temp_path, mode) as f:
                        for chunk in r.iter_content(chunk_size=self.chunk_size):
                            if chunk:  # 过滤keep-alive空块
                                f.write(chunk)
                        print('📦 模型压缩包下载完成！')

                # 原子操作：临时文件转正式文件
                self.temp_path.rename(self.zip_path)
                
                # 解压文件
                return self._extract_zip()

            except (requests.RequestException, IOError) as e:
                print(f'🔄 尝试 {attempt + 1}/{self.max_retries} 失败: {e}')
                if attempt < self.max_retries - 1:
                    time.sleep(2 ** attempt)  # 指数退避等待
        return ""  # 所有重试失败


if __name__ == "__main__":
    # 使用MinIO的压缩包进行测试
    minio_url = "http://localhost:9001/api/v1/download-shared-object/aHR0cDovLzEyNy4wLjAuMTo5MDAwL3pobjY4MTgvc2F2ZV9tb2RlbC56aXA_WC1BbXotQWxnb3JpdGhtPUFXUzQtSE1BQy1TSEEyNTYmWC1BbXotQ3JlZGVudGlhbD1URVIzU0FPTUxZS1dJWldYRzBaNiUyRjIwMjUxMDE2JTJGdXMtZWFzdC0xJTJGczMlMkZhd3M0X3JlcXVlc3QmWC1BbXotRGF0ZT0yMDI1MTAxNlQwMjQzNTVaJlgtQW16LUV4cGlyZXM9NDMyMDAmWC1BbXotU2VjdXJpdHktVG9rZW49ZXlKaGJHY2lPaUpJVXpVeE1pSXNJblI1Y0NJNklrcFhWQ0o5LmV5SmhZMk5sYzNOTFpYa2lPaUpVUlZJelUwRlBUVXhaUzFkSldsZFlSekJhTmlJc0ltVjRjQ0k2TVRjMk1EWXlNakkxT0N3aWNHRnlaVzUwSWpvaVlXUnRhVzRpZlEuN2VfSUtJNXY4VUxuUnpZSGRQcHZrdk5YaElsLTFqQW13czZLNTRrUHdpVjlKMmtpMnBUcXdNQ1psdkxNZU9ndUdPZzFUUUJLV0hYOWRmaHg0ZG1IbVEmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JnZlcnNpb25JZD1udWxsJlgtQW16LVNpZ25hdHVyZT02OTNkMmY0YWRmMjUxZDY5MGJmM2I5MGZiMmJkMmUxZGM0NmMxZGI4NDkxNjU3MGJlNGE4MmI5NDIwNmY0YWNk"
    res = ModelDownloader(minio_url, 'deploy_shuzhi').download
    print(f'模型压缩包下载地址:\n{res}')
